import yfinance as yf
import pandas as pd
import numpy as np
from scipy.optimize import minimize

_RISK_FREE_RATE = 0.03

# Download AAPL stock data and options data
ticker = "AAPL"
stock = yf.Ticker(ticker)

# Define the expiration date for the options (example: closest expiry date)
expiry_date = stock.options[0]

# Get the options chain for the given expiry date
options_chain = stock.option_chain(expiry_date)

# Filter out deep out-of-the-money (OTM) options
# For example, if AAPL is trading at $150, consider calls with strike > $180 and puts with strike < $120
spot_price = stock.history(period="1d")["Close"].iloc[-1]
deep_otm_calls = options_chain.calls[options_chain.calls["strike"] > spot_price * 1.2]
deep_otm_puts = options_chain.puts[options_chain.puts["strike"] < spot_price * 0.8]

# Combine and preview the data
deep_otm_options = pd.concat([deep_otm_calls, deep_otm_puts])
deep_otm_options["expiry"] = expiry_date

def simulate_merton_jump(
    init_price: float,
    expiry: int,
    mu: float,
    sigma: float,
    lambda_jump: float,
    jump_mean: float,
    jump_vol: float,
    num_paths: int = 100,
    num_steps: int = 252,
) -> np.ndarray:
    """
    Simulate Merton jump diffusion paths using Monte Carlo simulation.

    For instance, 30 days (physical date) rather than trading days to expiry, the T = 30/365.

    Args:
        S0: initial stock price.
        T: time to maturity.
        mu: drift term.
        sigma: diffusion term.
        lambda_jump: jump intensity, the average number of jumps per year.
        jump_mean: mean of jump size distribution.
        jump_std: standard deviation of jump size distribution.
        num_paths: number of simulated paths.
        num_steps: number of time steps, default to number of trading days in a year.
    """
    dt = expiry / num_steps
    paths = np.zeros(shape=(num_paths, num_steps + 1))
    paths[:, 0] = init_price  # Initial stock price for all paths

    for t in range(1, num_steps + 1):
        W = np.random.normal(loc=0, scale=1, size=num_paths)  # Standard Wiener process
        paths[:, t] = paths[:, t - 1] * np.exp(
            (mu - 0.5 * sigma**2) * dt + sigma * np.sqrt(dt) * W
        )
        num_jumps = np.random.poisson(lam=lambda_jump * dt, size=num_paths)
        jump_sizes = np.exp(jump_mean + jump_vol * np.random.normal(0, 1, num_paths))
        paths[:, t] *= np.where(num_jumps > 0, jump_sizes, 1)

    return paths

def merton_pricing(
    *,
    S0: float,
    K: float,
    T: int,
    r: float,
    mu: float,
    sigma: float,
    lambda_jump: float,
    jump_mean: float,
    jump_vol: float,
    num_paths,
    option_type: str,
) -> float:
    """
    Price European options under Merton jump diffusion model using simulated path.

    First generate simulated paths using the Merton jump diffusion model,
    then calculate the option price using the risk-neutral pricing formula.

    Args:
        S0: initial stock price.
        K: strike price.
        T: time to maturity.
        r: risk-free rate.
        mu: drift term.
        sigma: diffusion term.
        lambda_jump: jump intensity, the average number of jumps per year.
        jump_mean: mean of jump size distribution.
        jump_std: standard deviation of jump size distribution.
        num_paths: number of simulated paths.
        option_type: "call" or "put".

    Returns:
        BS option price.
    """
    paths = simulate_merton_jump(
        init_price=S0,
        expiry=T,
        mu=mu,
        sigma=sigma,
        lambda_jump=lambda_jump,
        jump_mean=jump_mean,
        jump_vol=jump_vol,
        num_paths=num_paths,
    )

    terminal_prices = paths[:, -1]  # Extract terminal prices

    if option_type == "call":
        payoffs = np.maximum(terminal_prices - K, 0)
    else:
        payoffs = np.maximum(K - terminal_prices, 0)

    return np.exp(-r * T) * np.mean(payoffs)  # risk neutral pricing

def objective_func(params: list, market_data, init_value, risk_free_rate) -> float:
    """
    Objective function to minimize the sum of squared
    errors between market prices and model prices.

    Args:
        params: model parameters, [mu, sigma, lambda_jump, jump_mean, jump_std].
        market_data: option chain data.
        init_value: initial stock price.
        risk_free_rate: risk-free rate.
    Returns:
        Sum of squared errors.
    """
    mu, sigma, lambda_jump, jump_mean, jump_std = params
    error = 0
    for _, row in market_data.iterrows():
        strike = row["strike"]
        expiry = (pd.to_datetime(row["expiry"]) - pd.Timestamp.now()).days / 365
        market_price = (row["bid"] + row["ask"]) / 2
        model_price = merton_pricing(
            S0=init_value,
            K=strike,
            T=expiry,
            r=risk_free_rate,
            mu=mu,
            sigma=sigma,
            lambda_jump=lambda_jump,
            jump_mean=jump_mean,
            jump_vol=jump_std,
            num_paths=100,
            option_type="call",
        )
        error += (market_price - model_price) ** 2
    return error

# Initial guesses for parameters
initial_params = [
    0.05,  # mu
    0.2,  # sigma
    0.1,  # lambda_jump
    -0.1,  # jump_mean
    0.1,  # jump_std
]

# Calibrate using deep OTM options
result = minimize(
    objective_func,
    initial_params,
    args=(deep_otm_options, spot_price, _RISK_FREE_RATE),
    method="Nelder-Mead",
)
calibrated_params = result.x
print("Calibrated parameters:", calibrated_params)

Calibrated parameters: [ 0.05200766  0.20268906  0.100778   -0.10256474  0.09463834]